import enum
import json
from shapely.geometry import Polygon, MultiPolygon
import os
from typing import Optional, List, Dict

from common.logger import logger

GEOJSON_PATHS = os.path.join(
    os.path.dirname(os.path.realpath(__file__)), "GeojsonPaths.json"
)
with open(GEOJSON_PATHS, "r") as f:
    geojson_dict = json.load(f)

SENTINEL2_GRID_GEOJSON = geojson_dict["SENTINEL2_GRID_GEOJSON"]
US_GEOJSON_COUNTY = geojson_dict["US_GEOJSON_COUNTY"]
US_GEOJSON_STATE = geojson_dict["US_GEOJSON_STATE"]
CHINA_GEOJSON = geojson_dict["CHINA_GEOJSON"]


class PlaceType(str, enum.Enum):
    STATE = "state"
    PREFECTURE = "prefecture"
    COUNTY = "county"


class GeoJSONPlace(object):
    """A Place service backed by data stored in geojson file(s)."""

    def __init__(self, path_to_json_file: str = None):
        """Initialize a `GeoJSONPlace` client connecting to a given backend json file

        Parameters
        ----------
        path_to_json_file : str, optional
            File path to a json file defining the backend data source,
            by default None to use predefined defaults.
        """
        # provide a default but also allow users to init with a different path
        # for easier development, testing, and upgrade
        s2_grid_geojson = SENTINEL2_GRID_GEOJSON
        us_county_path = US_GEOJSON_COUNTY
        us_state_path = US_GEOJSON_STATE
        china_geojson_path = CHINA_GEOJSON
        if path_to_json_file is not None:
            with open(path_to_json_file, "r") as f:
                geojson_dict = json.load(f)

            s2_grid_geojson = geojson_dict["SENTINEL2_GRID_GEOJSON"]
            us_county_path = geojson_dict["US_GEOJSON_COUNTY"]
            us_state_path = geojson_dict["US_GEOJSON_STATE"]
            china_geojson_path = geojson_dict["CHINA_GEOJSON"]

        self.Geojson_CN = GeoJSONPlace_CN(china_geojson_path, s2_grid_geojson)
        self.Geojson_US = GeoJSONPlace_US(
            us_state_path, us_county_path, s2_grid_geojson
        )

    def find_region(self, **kwargs):
        """Find region by given criteria

        Parameters
        ----------
        country: str
            Country name, support two options "CN" or "US".
        state_name: str
            State (in US) or province (in CN) name to be queried.
        county_name: str, optional
            County name to be queried.

        Returns
        -------
        List[Dict]
            List of found regions, each region as a dictionary following GeoJSON's `Feature` standard.

        Raises
        ------
        ValueError
            Raised when a `country` other than "US" or "CN" is given.
        """
        country = kwargs["country"]
        state_name = kwargs["state_name"]
        county_name = kwargs["county_name"] if "county_name" in kwargs else None
        if country.upper() == "CN":
            return self.Geojson_CN.get_feature_CN(state_name, county_name)
        elif country.upper() == "US":
            return self.Geojson_US.get_feature_US(state_name, county_name)
        else:
            raise ValueError("Input country of {} is not supported.".format(country))

    def find_mgrs_tiles(self, **kwargs):
        """Find MGRS tiles interesecting with given region specified by criteria.

        Parameters
        ----------
        country: str
            Country name, support two options "CN" or "US".
        state_name: str
            State (in US) or province (in CN) name to be queried.
        county_name: str, optional
            County name to be queried.

        Returns
        -------
        List[str]
            List of found MGRS tiles, each tile represented by a string uid,
            e.g. "T48RVV".

        Raises
        ------
        ValueError
            Raised when a `country` other than "US" or "CN" is given.
        """
        country = kwargs["country"]
        state_name = kwargs["state_name"]
        county_name = kwargs["county_name"] if "county_name" in kwargs else None
        if country.upper() == "CN":
            return self.Geojson_CN.find_tileid_CN(state_name, county_name)
        elif country.upper() == "US":
            return self.Geojson_US.find_tileid_US(state_name, county_name)
        else:
            raise ValueError("Input country of {} is not supported.".format(country))

    def find_place_ori(
        self, place_type, limit=10, country="CN", region_name=None, aoi=None, *kwargs
    ):
        """

        Parameters
        ----------
        place_type : [type]
            [description]
        limit : int, optional
            [description], by default 10
        country : str, optional
            [description], by default "CN"
        region_name : [type], optional
            [description], by default None
        aoi : [type], optional
            [description], by default None

        Raises
        ------
        NotImplementedError
            [description]

        Examples
        --------
        GeoJSONPlace.find_place("county", region_name="da4yi4")
        GeoJSONPlace.find_place(PlaceType.PREFECTURE, region_name="xiangyang", aoi=subset_of_xiangyang)
        """
        # reserved for future Point type of places
        raise NotImplementedError

    def find_region_ori(
        self,
        region_type: PlaceType,
        limit: int = 10,
        country: str = "CN",
        region_name: Optional[str] = None,
        aoi=None,
    ) -> List[Dict]:
        # find a region or regions by given criteria
        # region_type -> state/province, prefecture, county
        # country -> CN or CONUS for now
        # aoi -> let us discuss what to be supported
        #   for now, maybe, geojson Feature or FeatureCollection
        # for implementation, for efficiency maybe
        # - change the geojson files to sqlite geospatial database with index
        # - use gdal.VectorTranslate
        # return result as a List of Dictionaries
        # with each Dictionary as a representation of a geojson Feature
        # example
        # {"type":"Feature","properties": {<omitted>}, "geometry": {<omitted>}}
        pass

    def find_mgrs_tiles_ori(
        self, place_type, limit=10, country="CN", region_name=None, aoi=None
    ):
        # similar interface as others, but giving back MGRS tiles
        pass


class GeoJSONPlaceCountry(object):
    def __init__(self, s2_grid_geojson=SENTINEL2_GRID_GEOJSON) -> None:
        super().__init__()

        with open(s2_grid_geojson) as f:
            self.s2_geojson = json.load(f)

    def _geometry_to_tileid(self, geometry):
        if geometry is None:
            logger.warning("Current geometry input is None, return empty list")
            return []
        if geometry["type"] not in {"MultiPolygon", "Polygon"}:
            raise ValueError(
                "Geometry type should be MultiPolygon or Polygon, but is"
                + geometry["type"],
            )
        if geometry["type"] == "MultiPolygon":
            coord_list = []
            for coord in geometry["coordinates"]:
                coord_list.append(Polygon(coord[0]))
            target_polygon = MultiPolygon(coord_list)
        elif geometry["type"] == "Polygon":
            target_polygon = Polygon(geometry["coordinates"][0])

        else:
            logger.error("{} is not a supported geometry type".format(geometry["type"]))

            raise NotImplementedError

        tileid_set = set()

        for feature in self.s2_geojson["features"]:
            s2_polygon = Polygon(feature["geometry"]["coordinates"][0])
            if s2_polygon.intersects(target_polygon):
                tileid_set.add(feature["properties"]["Name"])

        return tileid_set

    def get_tileid_with_geojson(self, target_geojson):
        if not isinstance(target_geojson, (dict, str)):
            raise ValueError(
                "Not a valid type for input geojson, should either be geojson dict or geojson file path"
            )
        max_area_num = 60
        if isinstance(target_geojson, str):
            with open(
                target_geojson, "r"
            ) as f:  # if cannot open try encoding="latin-1"
                target = json.load(f)

        elif isinstance(target_geojson, dict):
            target = target_geojson

        tileid_set = set()

        if len(target["features"]) > max_area_num:
            logger.error(
                "Too many shapes in this geojson, No more than {} shape is accepted.".format(
                    max_area_num
                )
            )
            return []

        for feature in target["features"]:
            geometry = feature["geometry"]
            tmp_tile_set = self._geometry_to_tileid(geometry)
            for val in tmp_tile_set:
                tileid_set.add(val)

        tileid_list = list(tileid_set)
        tileid_list.sort()
        return tileid_list


class GeoJSONPlace_CN(GeoJSONPlaceCountry):
    def __init__(
        self, geojson_root=CHINA_GEOJSON, s2_grid_geojson=SENTINEL2_GRID_GEOJSON
    ) -> None:
        super().__init__(s2_grid_geojson)
        self.china_json_root = geojson_root

        self.province_path = [os.path.join(self.china_json_root, "省")]
        self.city_path = [
            os.path.join(self.china_json_root, "市"),
            os.path.join(self.china_json_root, "县"),
        ]

    def _find_geojson_path_CN(self, province_name, county_name=None):
        filepath_list = []

        if county_name is None:
            path_list = self.province_path
            target_name = province_name

        else:
            path_list = self.city_path
            target_name = county_name

        for path_root in path_list:
            for root, _, filenames in os.walk(path_root):
                for basename in filenames:
                    if target_name in basename:
                        file_path = os.path.join(root, basename)
                        filepath_list.append(file_path)

        filepath_list.sort()
        return filepath_list

    def __get_json_dict(self, json_path):
        with open(json_path, encoding="latin-1") as f:
            json_dict = json.load(f)

        return json_dict

    def get_geometry_CN(self, province_name, county_name=None):
        json_path_list = self._find_geojson_path_CN(province_name, county_name)
        if not json_path_list:
            logger.warning(
                "No json path found for area in CN with name {} {}\nReturning None".format(
                    province_name, county_name
                )
            )
            return None

        elif len(json_path_list) > 1:
            json_name_list = []
            for filepath in json_path_list:
                basename = os.path.basename(filepath)
                json_name_list.append(basename)
            logger.warning(
                "More than 1 area found:\n{}\nReturning None".format(json_name_list)
            )
            return None

        else:
            json_dict = self.__get_json_dict(json_path_list[0])
            geometry = None
            for feature in json_dict["features"]:
                geometry = feature["geometry"]

            if geometry is None:
                logger.warning(
                    "Json containing no geometry.\n{}".format(json_path_list[0])
                )

            return geometry

    def find_tileid_CN(self, province_name, county_name=None):
        geometry = self.get_geometry_CN(province_name, county_name)
        tileid_list = list(self._geometry_to_tileid(geometry))
        tileid_list.sort()
        return tileid_list

    def get_feature_CN(self, province_name, county_name=None):
        json_path_list = self._find_geojson_path_CN(province_name, county_name)
        if not json_path_list:
            logger.warning(
                "No json path found for area in CN with name {} {}\nReturning None".format(
                    province_name, county_name
                )
            )
            return None

        elif len(json_path_list) > 1:
            json_name_list = []
            for filepath in json_path_list:
                basename = os.path.basename(filepath)
                json_name_list.append(basename)
            logger.warning(
                "More than 1 area found:\n{}\nReturning None".format(json_name_list)
            )
            return None

        else:
            json_dict = self.__get_json_dict(json_path_list[0])
            feature_list = []
            for feature in json_dict["features"]:
                feature_list.append(feature)

            if not feature_list:
                logger.warning(
                    "Json containing no features.\n{}".format(json_path_list[0])
                )

            return feature_list


class GeoJSONPlace_US(GeoJSONPlaceCountry):
    def __init__(
        self,
        us_geojson_state_path=US_GEOJSON_STATE,
        us_geojson_county_path=US_GEOJSON_COUNTY,
        s2_grid_geojson=SENTINEL2_GRID_GEOJSON,
    ) -> None:
        super().__init__(s2_grid_geojson)

        with open(us_geojson_state_path, encoding="latin-1") as f:
            self.us_geojson_state = json.load(f)
        self.us_state_to_id_dict = dict()
        for feature in self.us_geojson_state["features"]:
            self.us_state_to_id_dict[feature["properties"]["NAME"]] = feature[
                "properties"
            ]["STATE"]

        with open(us_geojson_county_path, encoding="latin-1") as f:
            self.us_geojson_county = json.load(f)

    def __get_feature_US_STATE(self, state_name):
        feature_list = []
        for feature in self.us_geojson_state["features"]:
            if feature["properties"]["NAME"] == state_name:
                feature_list.append(feature)

        if not feature_list:
            logger.warning('No feature for State: "{}" found.'.format(state_name))

        return feature_list

    def __get_feature_US_COUNTY(self, state_name, county_name):
        feature_list = []
        state_id = self.us_state_to_id_dict[state_name]
        for feature in self.us_geojson_county["features"]:
            if (
                feature["properties"]["NAME"] == county_name
                and feature["properties"]["STATE"] == state_id
            ):
                feature_list.append(feature)

        if not feature_list:
            logger.warning(
                'No feature for State: "{} {}" found.'.format(state_name, county_name)
            )

        return feature_list

    def __get_geometry_US_STATE(self, state_name):
        geometry = None
        for feature in self.us_geojson_state["features"]:
            if feature["properties"]["NAME"] == state_name:
                geometry = feature["geometry"]

        if geometry is None:
            logger.warning('No geometry for State: "{}" found.'.format(state_name))

        return geometry

    def __get_geometry_US_COUNTY(self, state_name, county_name):
        if state_name not in self.us_state_to_id_dict.keys():
            raise ValueError("{} is not a valid state name".format(state_name))
        state_id = self.us_state_to_id_dict[state_name]
        geometry = None
        for feature in self.us_geojson_county["features"]:
            if (
                feature["properties"]["NAME"] == county_name
                and feature["properties"]["STATE"] == state_id
            ):
                geometry = feature["geometry"]

        if geometry is None:
            logger.warning(
                'No County: "{}" found in State: "{}"'.format(county_name, state_name)
            )
        return geometry

    def get_geometry_US(self, state_name, county_name=None):
        geometry = None
        if county_name is None:
            geometry = self.__get_geometry_US_STATE(state_name)
        else:
            geometry = self.__get_geometry_US_COUNTY(state_name, county_name)

        return geometry

    def find_tileid_US(self, state_name, county_name=None):
        tileid_list = []
        geometry = self.get_geometry_US(state_name, county_name)
        tileid_list = list(self._geometry_to_tileid(geometry))
        tileid_list.sort()
        return tileid_list

    def get_feature_US(self, state_name, county_name=None):
        feature_list = []
        if county_name is None:
            feature_list = self.__get_feature_US_STATE(state_name)
        else:
            feature_list = self.__get_feature_US_COUNTY(state_name, county_name)

        return feature_list


if __name__ == "__main__":
    helper = GeoJSONPlace()
    tile_list = helper.find_mgrs_tiles(
        country="US", state_name="California", county_name=None
    )
    tile_list2 = helper.find_mgrs_tiles(country="CN", state_name="", county_name="平邑")
    tile_list3 = helper.find_mgrs_tiles(country="CN", state_name="", county_name="临沭")
    feature1 = helper.find_region(
        country="US", state_name="California", county_name="Santa Clara"
    )
    feature2 = helper.find_region(country="CN", state_name="", county_name="临沭")
    print(tile_list2)
    print(tile_list3)
